//! Crank-Nicholson 方法求解偏微分方程
# ifndef CrankNicholson_H_
# define CrankNicholson_H_

#include <iostream>
#include "TDMA.h"
#include<complex>
#define complex std::complex<double>

//CrankNicholson--------------------------------------------------------------------------
class CrankNicholson {
private:
    int m_nr;//径向r 的格点数
    complex * d; //A*X=D
    TDMA* tdma;
public:
    
    CrankNicholson(int nr) :m_nr(nr){//初始化列表  ,m_nz(nz) 
        d = new complex[m_nr];
        tdma = new TDMA(m_nr);   
    }
    virtual ~CrankNicholson() {
        //cout <<"调用CN析构函数"<<endl;
        delete[] d;
        delete tdma;
    }

    //Get_d
    template<typename B0, typename B1, typename B2>
    void Get_d(B0 b0, B1 b1, B2 b2, complex* X0, complex* G0);
    //Get_X
    template<typename A0, typename A1, typename A2>
    void Get_x(A0 a0, A1 a1, A2 a2, complex* X1);

    template<typename A0, typename A1, typename A2, typename B0, typename B1, typename B2>
    void operator()(A0 a0, A1 a1, A2 a2, 
                    B0 b0, B1 b1, B2 b2, 
                    complex* X0, complex* G0, complex* X1);
};

//Get_d
template<typename B0, typename B1, typename B2>
void CrankNicholson::Get_d(B0 b0, B1 b1, B2 b2, complex* X0, complex* G0){ //求解频域传播方程，输入频域电场数据X0，和频域非线性项G0


    d[0]=b1[0] * X0[0] + b2[0] * X0[1] - G0[0];
    for(int id = 1; id < m_nr-1; id++){
        d[id]=b0[id-1] * X0[id-1] + b1[id] * X0[id] + b2[id] * X0[id+1] - G0[id]; 
    }

    d[m_nr-1] = b0[m_nr-2] * X0[m_nr-2] + b1[m_nr-1] * X0[m_nr-1] - G0[m_nr -1];  

}

//Get_X1
template<typename A0, typename A1, typename A2>
void CrankNicholson::Get_x(A0 a0, A1 a1, A2 a2, complex* X1) {

    tdma->solve(a0, a1, a2, d, X1);
    
}

template<typename A0, typename A1, typename A2, typename B0, typename B1, typename B2>
void CrankNicholson::operator()
                (A0 a0, A1 a1, A2 a2, 
                B0 b0, B1 b1, B2 b2, 
                complex* X0, complex* G0, complex* X1){
    Get_d(b0, b1, b2, X0, G0);
    tdma->solve(a0, a1, a2, d, X1);
}

# endif








